{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Catalog Search Agent - Find Datasets with Natural Language\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/catalog_search_agent.ipynb)\n",
    "\n",
    "This notebook demonstrates how to use the Catalog Search Agent to find datasets in data catalogs using natural language queries.\n",
    "\n",
    "The Catalog Agent can:\n",
    "- Search through data catalogs (TSV, CSV, JSON formats)\n",
    "- **Search by geographic region** (location names or bounding boxes)\n",
    "- Understand natural language queries\n",
    "- Filter by keywords, dataset type, and provider\n",
    "- Return structured results with metadata\n",
    "- Work with Earth Engine Data Catalog and custom catalogs\n",
    "\n",
    "Uncomment the following line to install [geoai](https://opengeoai.org) if needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install \"geoai-py[agents]\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from geoai.agents import (\n",
    "    CatalogAgent,\n",
    "    CatalogTools,\n",
    "    create_ollama_model,\n",
    "    create_openai_model,\n",
    "    create_anthropic_model,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Earth Engine Data Catalog\n",
    "\n",
    "The Earth Engine Data Catalog is available in **JSON format** from the [opengeos/Earth-Engine-Catalog](https://github.com/opengeos/Earth-Engine-Catalog) repository.\n",
    "\n",
    "**Important:** Use the JSON format (not TSV) to enable spatial search capabilities with bounding box information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catalog_url = \"https://raw.githubusercontent.com/opengeos/Earth-Engine-Catalog/refs/heads/master/gee_catalog.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a model\n",
    "\n",
    "You can create a model with the following functions:\n",
    "\n",
    "- `create_ollama_model`: Create a model using Ollama. You will need to install [Ollama](https://ollama.com/download) separately and pull the model you want to use, such as `llama3.1`.\n",
    "- `create_openai_model`: Create a model using OpenAI. You will need an OpenAI API key. Set it in the `OPENAI_API_KEY` environment variable.\n",
    "- `create_anthropic_model`: Create a model using Anthropic. You will need an Anthropic API key. Set it in the `ANTHROPIC_API_KEY` environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_ollama_model(model=\"llama3.1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use CatalogTools Directly (Fast)\n",
    "\n",
    "For faster searches without LLM overhead, you can use `CatalogTools` directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = CatalogTools(catalog_url=catalog_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get catalog statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = json.loads(tools.get_catalog_stats())\n",
    "print(f\"Total datasets: {stats['total_datasets']}\")\n",
    "print(f\"\\nDataset types:\")\n",
    "for dtype, count in stats.get(\"dataset_types\", {}).items():\n",
    "    print(f\"  {dtype}: {count}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Search for datasets by keyword"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = json.loads(tools.search_datasets(keywords=\"landcover\", max_results=5))\n",
    "print(f\"Found {result['dataset_count']} datasets\\n\")\n",
    "\n",
    "for ds in result[\"datasets\"]:\n",
    "    print(f\"ID: {ds['id']}\")\n",
    "    print(f\"Title: {ds['title']}\")\n",
    "    print(f\"Provider: {ds.get('provider', 'N/A')}\")\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Search by geographic region (NEW!)\n",
    "\n",
    "Find datasets covering a specific location or bounding box:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search by location name\n",
    "result = json.loads(\n",
    "    tools.search_by_region(location=\"California\", keywords=\"elevation\", max_results=5)\n",
    ")\n",
    "\n",
    "print(f\"Found {result['dataset_count']} datasets covering California\\n\")\n",
    "\n",
    "for ds in result[\"datasets\"]:\n",
    "    print(f\"ID: {ds['id']}\")\n",
    "    print(f\"Title: {ds['title']}\")\n",
    "    print(f\"Bbox: {ds.get('bbox', 'N/A')}\")\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search by bounding box coordinates\n",
    "# San Francisco Bay Area: [west, south, east, north]\n",
    "result = json.loads(\n",
    "    tools.search_by_region(\n",
    "        bbox=[-122.5, 37.5, -122.0, 38.0], keywords=\"landcover\", max_results=3\n",
    "    )\n",
    ")\n",
    "\n",
    "print(f\"Found {result['dataset_count']} datasets\\n\")\n",
    "\n",
    "for ds in result[\"datasets\"]:\n",
    "    print(f\"ID: {ds['id']}\")\n",
    "    print(f\"Title: {ds['title']}\")\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Geocode location names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert location name to bounding box\n",
    "location_info = json.loads(tools.geocode_location(\"New York City\"))\n",
    "print(f\"Name: {location_info['name']}\")\n",
    "print(f\"Bbox: {location_info['bbox']}\")\n",
    "print(f\"Center: {location_info['center']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Search with filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = json.loads(\n",
    "    tools.search_datasets(\n",
    "        keywords=\"elevation\", dataset_type=\"image\", provider=\"NASA\", max_results=5\n",
    "    )\n",
    ")\n",
    "\n",
    "print(f\"Found {result['dataset_count']} datasets\\n\")\n",
    "\n",
    "for ds in result[\"datasets\"]:\n",
    "    print(f\"ID: {ds['id']}\")\n",
    "    print(f\"Title: {ds['title']}\")\n",
    "    print(f\"Type: {ds.get('type', 'N/A')}\")\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List available providers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "providers = json.loads(tools.list_providers())\n",
    "print(f\"Total providers: {providers['count']}\\n\")\n",
    "print(\"Sample providers:\")\n",
    "for p in providers[\"providers\"][:10]:\n",
    "    print(f\"  - {p}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List dataset types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "types = json.loads(tools.list_dataset_types())\n",
    "print(f\"Available dataset types: {types['types']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use CatalogAgent with Natural Language (LLM)\n",
    "\n",
    "For natural language queries, create a `CatalogAgent` that uses an LLM to understand and execute searches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = CatalogAgent(model=model, catalog_url=catalog_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ask natural language questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Find datasets about landcover from NASA\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Show me elevation data\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial search with natural language (NEW!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Find landcover datasets covering California\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Show me elevation data for San Francisco\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Find land cover datasets from NASA that cover New York City\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"What types of datasets are available?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Find image collections about forests\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.ask(\"Find landcover datasets from 2022 onwards\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get structured results programmatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = agent.search_datasets(keywords=\"sentinel\", max_results=5)\n",
    "\n",
    "for ds in datasets:\n",
    "    print(f\"ID: {ds['id']}\")\n",
    "    print(f\"Title: {ds['title']}\")\n",
    "    print(f\"Provider: {ds.get('provider', 'N/A')}\")\n",
    "    if ds.get(\"snippet\"):\n",
    "        print(f\"Code: {ds['snippet']}\")\n",
    "    print(\"-\" * 80)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
